package cryptosquare

import (
	"math"
	"regexp"
	"strings"
)

const testVersion = 2

var regLetters = regexp.MustCompile(`[a-z0-9]+`)

func Encode(input string) string {
	s1 := strings.Join(regLetters.FindAllString(strings.ToLower(input), -1), "")
	newS := getNewInputString(s1)

	l, c, r := getColumnsAndRows(s1)
	newSArray := make([]string, len(newS)+r+1)
	fix := (len(newS)%r + 2) * c
	step := 0
	if c == r {
		for i := 0; i < len(newSArray); i++ {
			if (i+1)%(r+1) == 0 {
				newSArray[i] = " "
				continue
			}
			newSArray[i] = newS[step : step+1]
			step++
			if step >= l {
				break
			}
		}
	} else {
		for i := 0; i < len(newSArray); i++ {
			if i > 0 && (i+1)%(r+1) == 0 {
				if i < fix {
					newSArray[i] = " "
				} else {
					newSArray[i] = " "
				}
				continue
			}

			if i > fix && i%c == c-1 {
				newSArray[i] = " "
				continue
			}

			newSArray[i] = newS[step : step+1]
			step++
			if step >= l {
				break
			}
		}
	}

	return strings.Join(newSArray, "")
}
func getColumnsAndRows(input string) (inputLength, columns, rows int) {

	inputLength = len(input)
	rows = int(math.Pow(float64(inputLength), 0.5))
	if rows*rows == inputLength {
		columns = rows
	} else {
		columns = rows + 1
	}

	return
}

func getNewInputString(input string) (output string) {
	l, c, r := getColumnsAndRows(input)
	sArray := make([]string, c*c)
	idx := 0
	for i := 0; i <= r; i++ {
		for j := 0; j < c; j++ {
			if j*c+i >= l {
				continue
			}
			if i*r+j >= l {
				break
			}

			sArray[idx] = input[j*c+i : j*c+i+1]
			idx++
		}
	}

	output = strings.Join(sArray, "")
	return
}
